import sys
import os
import yaml
import markdown
import logging
import time
import tiktoken
from datetime import datetime
from dotenv import load_dotenv
from openai import OpenAI
from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout, 
                           QTextEdit, QPushButton, QComboBox, QSplitter, QLabel)
from PyQt5.QtCore import Qt, QThread, pyqtSignal

# 设置日志格式
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

# 加载环境变量
load_dotenv()

def load_config():
    config_path = os.path.join(os.path.dirname(__file__), 'tools.yaml')
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
    except Exception as e:
        print(f"Error loading config: {e}")
        return {"name": "InChat", "title": {"en_US": "InChat"}, "version": {"en_US": "v1"}}

# 创建自定义日志处理器，用于GUI显示
class QTextEditLogger(logging.Handler):
    """用于将日志输出到QTextEdit控件的自定义日志处理器类
    
    主要功能:
    - 接收日志记录并显示到指定的QTextEdit控件中
    - 支持自定义日志格式
    """
    def __init__(self, widget):
        """初始化日志处理器
        
        Args:
            widget: QTextEdit控件实例，用于显示日志内容
        """
        super().__init__()
        self.widget = widget
        self.widget.setReadOnly(True)
        
    def emit(self, record):
        """输出日志记录到QTextEdit控件
        
        Args:
            record: 日志记录对象，包含日志的详细信息
        """
        msg = self.format(record)
        self.widget.append(msg)

def count_tokens(text, model="gpt-3.5-turbo"):
    """使用tiktoken计算文本的token数量
    
    Args:
        text: 要计算的文本内容
        model: 使用的模型名称
    
    Returns:
        int: token数量
    """
    try:
        encoding = tiktoken.encoding_for_model(model)
    except KeyError:
        # 如果模型不在预设列表中，使用cl100k_base编码
        encoding = tiktoken.get_encoding("cl100k_base")
    return len(encoding.encode(text))

# AIWorker类
class AIWorker(QThread):
    """AI对话处理线程类
    
    主要功能:
    - 在后台线程中处理与AI的对话
    - 使用OpenAI兼容接口进行流式对话
    - 通过信号机制实时返回AI的响应内容
    """
    
    update_signal = pyqtSignal(str)
    # 添加新的信号用于传递token统计信息
    tokens_signal = pyqtSignal(dict)

    def __init__(self, user_input, config):
        """初始化AI工作线程
        
        Args:
            user_input: 用户输入的文本内容
            config: 配置信息字典
        """
        super().__init__()
        self.user_input = user_input
        self.config = config
        
        # 初始化 OpenAI 客户端
        self.client = OpenAI(
            base_url=os.getenv('OPENAI_API_BASE'),
            api_key=os.getenv('OPENAI_API_KEY')
        )

    def run(self):
        """线程执行方法，处理AI对话请求
        
        功能:
        - 调用OpenAI接口进行对话
        - 处理流式响应
        - 发送响应内容更新信号
        """
        try:
            response = None
            # 使用tiktoken计算prompt的tokens
            prompt_tokens = count_tokens(self.user_input)
            completion_tokens = 0
            accumulated_response = ""

            stream = self.client.chat.completions.create(
                model=os.getenv('OPENAI_API_MODEL', 'deepseek-chat'),
                messages=[{'role': 'user', 'content': self.user_input}],
                stream=True
            )

            for chunk in stream:
                if chunk.choices[0].delta.content is not None:
                    content = chunk.choices[0].delta.content
                    accumulated_response += content
                    self.update_signal.emit(content)
                # 保存最后一个响应对象
                response = chunk

            # 计算完成后的tokens
            completion_tokens = count_tokens(accumulated_response)
            
            # 发送token统计信息
            # 如果response中有usage信息就使用response中的，否则使用tiktoken计算的结果
            if response and hasattr(response, 'usage'):
                if response.usage is not None:
                    tokens_info = {
                        'total_tokens': response.usage.total_tokens,
                        'prompt_tokens': response.usage.prompt_tokens,
                        'completion_tokens': response.usage.completion_tokens
                    }
                else:
                    tokens_info = {
                        'total_tokens': prompt_tokens + completion_tokens,
                        'prompt_tokens': prompt_tokens,
                        'completion_tokens': completion_tokens
                    }
                self.tokens_signal.emit(tokens_info)

        except Exception as e:
            error_prefix = self.config['ui'].get('zh_CN', {}).get('error_prefix', '[Error]')
            self.update_signal.emit(f"\n{error_prefix} {str(e)}")


class ChatWindow(QWidget):
    """聊天窗口主类
    
    主要功能:
    - 提供聊天界面和交互功能
    - 支持多语言切换
    - 支持Markdown格式显示
    - 集成日志显示
    """
    
    def __init__(self):
        """初始化聊天窗口
        
        功能:
        - 初始化日志记录器
        - 加载配置信息
        - 初始化界面组件
        """
        super().__init__()
        # 创建日志记录器
        self.logger = logging.getLogger('InChat')
        self.logger.setLevel(logging.INFO)
        
        # 加载配置
        self.config = load_config()
        self.lang = 'zh_CN'  # 默认使用中文
        self.ui_text = self.config['ui'].get(self.lang, {})
        
        # 修正 Markdown 转换器初始化
        self.md = markdown.Markdown(
            extensions=['fenced_code', 'tables', 'nl2br', 'codehilite']
        )
        
        self.current_response = ""  # 用于累积流式响应
        self.last_cursor_position = 0  # 添加光标位置追踪
        self.ai_thread = None
        
        self.start_time = 0  # 用于计时
        self.token_count = 0  # 用于统计令牌数
        self.prompt_tokens = 0
        self.completion_tokens = 0
        
        self.log_visible = True  # 添加日志显示状态标记
        
        self.setStyleSheet("""
            QWidget {
                font-family: "Microsoft YaHei", "Segoe UI", Arial;
                font-size: 14px;
            }
            QComboBox {
                border: 1px solid #ddd;
                border-radius: 4px;
                padding: 5px;
                background: white;
                min-height: 25px;
            }
            QComboBox:hover {
                border-color: #0078d4;
            }
            QPushButton {
                border: none;
                border-radius: 4px;
                padding: 8px 16px;
                color: white;
                font-weight: 500;
            }
            QPushButton:hover {
                opacity: 0.9;
            }
            QTextEdit {
                border: 1px solid #ddd;
                border-radius: 4px;
                padding: 8px;
                background: white;
            }
        """)

        # 初始化UI必须在最后，因为它依赖于上面的所有属性
        self.init_ui()

    def init_ui(self):
        """初始化用户界面
        
        功能:
        - 创建并设置所有界面组件
        - 设置左右分栏布局
        - 配置信号和槽的连接
        """
        # 设置窗口标题
        self.update_window_title()
        self.setGeometry(100, 100, 1000, 700)  # 加大默认窗口尺寸

        # 创建主布局为水平布局
        main_layout = QHBoxLayout()
        main_layout.setSpacing(0)
        main_layout.setContentsMargins(0, 0, 0, 0)
        
        # 创建左侧面板
        left_panel = QWidget()
        left_panel.setFixedWidth(120)  # 固定宽度
        left_panel.setStyleSheet("""
            QWidget {
                background-color: #2c2c2c;
                color: white;
            }
            QPushButton {
                text-align: left;
                padding: 12px;
                border-radius: 0;
                width: 100%;
                margin: 0;
            }
            QPushButton:hover {
                background-color: #404040;
            }
            QComboBox {
                background-color: #404040;
                color: white;
                border: none;
                border-radius: 0;
                padding: 8px;
                margin: 4px 0;
            }
        """)
        left_layout = QVBoxLayout(left_panel)
        
        # 添加语言选择下拉框到左侧
        self.lang_selector = QComboBox()
        self.lang_selector.addItems(['简体中文', 'English'])
        self.lang_selector.setCurrentText('简体中文' if self.lang == 'zh_CN' else 'English')
        self.lang_selector.currentTextChanged.connect(self.change_language)
        left_layout.addWidget(self.lang_selector)
        
        # 在语言选择器后添加日志开关按钮
        self.toggle_log_btn = QPushButton(self.ui_text.get('toggle_log', 'Toggle Log'))
        self.toggle_log_btn.setStyleSheet("""
            QPushButton {
                background-color: #17a2b8;
                color: white;
                border: none;
                padding: 5px 15px;
                font-size: 14px;
                border-radius: 3px;
                margin-top: 10px;
            }
            QPushButton:hover { 
                background-color: #138496;
            }
        """)
        self.toggle_log_btn.clicked.connect(self.toggle_log_display)
        left_layout.addWidget(self.toggle_log_btn)
        
        # 添加空白占位
        left_layout.addStretch()
        
        # 添加退出按钮到左侧底部
        self.exit_btn = QPushButton(self.ui_text.get('exit_button', 'Exit'))
        self.exit_btn.setStyleSheet("""
            QPushButton {
                background-color: #dc3545;
                color: white;
                border: none;
                padding: 5px 15px;
                font-size: 14px;
                border-radius: 3px;
                margin-bottom: 10px;
            }
            QPushButton:hover { 
                background-color: #c82333;
            }
        """)
        self.exit_btn.clicked.connect(self.close)
        left_layout.addWidget(self.exit_btn)
        
        # 创建右侧面板
        right_panel = QWidget()
        right_panel.setStyleSheet("""
            QWidget {
                background-color: #f5f5f5;
            }
            QTextEdit {
                border: 1px solid #e0e0e0;
                border-radius: 4px;
                background-color: white;
            }
            QTextEdit#chatDisplay {
                font-size: 15px;
                line-height: 1.6;
            }
            QLabel {
                color: #666;
                background: white;
                border: 1px solid #e0e0e0;
                border-radius: 4px;
                padding: 8px;
                margin: 4px 0;
            }
        """)
        right_layout = QVBoxLayout(right_panel)
        
        # 创建垂直分割器用于聊天区域和日志区域
        splitter = QSplitter(Qt.Vertical)
        
        # 聊天区域
        chat_widget = QWidget()
        chat_layout = QVBoxLayout(chat_widget)
        chat_layout.setSpacing(10)  # 设置垂直间距为10
        
        # 聊天显示框
        self.chat_display = QTextEdit()
        self.chat_display.setReadOnly(True)
        self.chat_display.setAcceptRichText(True)
        self.chat_display.setObjectName("chatDisplay")
        self.chat_display.setStyleSheet("""
            QTextEdit {
                border: none;
                background-color: white;
                padding: 15px;
            }
            QTextEdit code {
                background-color: #f8f8f8;
                padding: 2px 4px;
                border-radius: 3px;
                font-family: "Consolas", "Monaco", monospace;
                font-size: 13px;
            }
        """)

        # 添加速率显示标签到聊天布局
        self.speed_label = QLabel()
        self.speed_label.setStyleSheet("""
            QLabel {
                color: #666;
                background-color: white;
                border: 1px solid #e0e0e0;
                border-radius: 4px;
                padding: 8px;
                margin: 4px 0;
                font-family: "Consolas", monospace;
            }
        """)

        self.speed_label.setVisible(False)  # 初始时隐藏
        
        # 输入区域
        input_widget = QWidget()
        input_layout = QHBoxLayout(input_widget)  # 改为水平布局
        input_layout.setContentsMargins(15, 10, 15, 15)
        input_layout.setSpacing(10)  # 设置水平间距为10
        
        # 输入框
        self.input_area = QTextEdit()
        self.input_area.setMaximumHeight(100)
        self.input_area.setPlaceholderText(self.ui_text.get('input_placeholder', 'Enter your question...'))
        self.input_area.installEventFilter(self)
        self.input_area.setStyleSheet("""
            QTextEdit {
                border: 1px solid #e0e0e0;
                border-radius: 4px;
                padding: 10px;
                font-size: 14px;
                line-height: 1.5;
            }
            QTextEdit:focus {
                border-color: #0078d4;
            }
        """)

        # 发送按钮
        self.send_btn = QPushButton(self.ui_text.get('send_button', 'Send'))
        self.send_btn.setFixedWidth(60)  # 固定按钮宽度
        self.send_btn.setStyleSheet("""
            QPushButton {
                background-color: #0078d4;
                color: white;
                border-radius: 4px;
                padding: 10px 20px;
                font-weight: 500;
                min-width: 80px;
            }
            QPushButton:hover {
                background-color: #106ebe;
            }
            QPushButton:pressed {
                background-color: #005a9e;
            }
        """)
        self.send_btn.clicked.connect(self.send_message)

        # 添加组件到输入布局
        input_layout.addWidget(self.input_area)
        input_layout.addWidget(self.send_btn)
        
        # 添加组件到聊天布局
        chat_layout.addWidget(self.chat_display)
        chat_layout.addWidget(self.speed_label)  # 在聊天显示框和输入区域之间添加速率标签
        chat_layout.addWidget(input_widget)
        
        # 创建日志显示区域
        self.log_display = QTextEdit()
        self.log_display.setStyleSheet("""
            QTextEdit {
                background-color: #1e1e1e;
                color: #d4d4d4;
                border: none;
                border-radius: 4px;
                padding: 10px;
                font-family: "Consolas", monospace;
                font-size: 12px;
            }
        """)
        
        # 设置日志处理器
        log_handler = QTextEditLogger(self.log_display)
        log_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
        self.logger.addHandler(log_handler)

        # 添加组件到分割器
        splitter.addWidget(chat_widget)
        splitter.addWidget(self.log_display)
        splitter.setStretchFactor(0, 7)  # 聊天区域占70%
        splitter.setStretchFactor(1, 3)  # 日志区域占30%
        
        # 将分割器添加到右侧布局
        right_layout.addWidget(splitter)
        
        # 将左右面板添加到主布局
        main_layout.addWidget(left_panel)
        main_layout.addWidget(right_panel)
        
        self.setLayout(main_layout)

    def update_window_title(self):
        """更新窗口标题
        
        功能:
        - 根据当前语言设置从配置中获取标题和版本号
        - 更新窗口标题显示
        """
        title = self.config['title'].get(self.lang, 'InChat')
        version = self.config['version'].get(self.lang, 'v1')
        self.setWindowTitle(f"{title} {version}")

    def change_language(self, lang_text):
        """切换界面语言
        
        Args:
            lang_text: 语言选项文本('简体中文'或'English')
            
        功能:
        - 更新当前语言设置
        - 重新加载界面文字
        - 更新所有界面组件的显示文本
        """
        # 根据选择更新当前语言
        self.lang = 'zh_CN' if lang_text == '简体中文' else 'en_US'
        self.ui_text = self.config['ui'].get(self.lang, {})
        
        # 更新界面文字
        self.update_window_title()
        self.input_area.setPlaceholderText(self.ui_text.get('input_placeholder', 'Enter your question...'))
        self.send_btn.setText(self.ui_text.get('send_button', 'Send'))
        self.exit_btn.setText(self.ui_text.get('exit_button', 'Exit'))
        self.toggle_log_btn.setText(self.ui_text.get('toggle_log', 'Toggle Log'))
        # 如果速率标签可见，更新其文本
        if self.speed_label.isVisible():
            elapsed_time = time.time() - self.start_time
            if elapsed_time > 0:
                tokens_per_sec = self.token_count / elapsed_time
                self.speed_label.setText(
                    self.ui_text.get('log_speed', 'Response speed: {:.2f} Tokens/Sec').format(tokens_per_sec)
                )

    def eventFilter(self, obj, event):
        """事件过滤器，处理快捷键事件
        
        Args:
            obj: 触发事件的对象
            event: 事件对象
            
        功能:
        - 处理Ctrl+Enter发送消息快捷键
        """
        if obj == self.input_area and event.type() == event.KeyPress:
            if event.key() == Qt.Key_Return and event.modifiers() == Qt.ControlModifier:
                self.send_message()
                return True
        return super().eventFilter(obj, event)

    def send_message(self):
        """发送消息处理函数
        
        功能:
        - 获取并处理用户输入
        - 在聊天窗口显示用户消息
        - 启动AI处理线程
        - 记录相关日志
        """
        user_input = self.input_area.toPlainText().strip()
        if not user_input:
            return

        # 记录用户输入，使用配置的提示语
        self.logger.info(f"{self.ui_text.get('log_user_input', 'User Input: ')}{user_input[:50]}...")
        
        # 在聊天窗口中显示用户输入的消息
        self._append_message(self.ui_text.get('you', 'You'), user_input)
        self.input_area.clear()  # 清空输入框

        # 在聊天窗口中显示“思考中...”提示
        self._append_message(self.ui_text.get('bot', 'Bot'), 
                           self.ui_text.get('thinking', 'Thinking...'), 
                           is_streaming=True)

        # 启动后台线程，处理与AI的对话
        self.ai_thread = AIWorker(user_input, self.config)
        # 连接信号和槽函数，更新聊天内容
        self.ai_thread.update_signal.connect(self.update_bot_response)
        # 连接token统计信号
        self.ai_thread.tokens_signal.connect(self.update_tokens_info)
        self.ai_thread.finished.connect(self.finalize_response)
        
        self.start_time = time.time()  # 开始计时
        self.token_count = 0  # 重置令牌计数
        self.speed_label.setVisible(False)  # 发送新消息时隐藏速率显示
        
        # 启动线程
        self.ai_thread.start()

    def update_tokens_info(self, tokens_info):
        """更新token统计信息"""
        self.prompt_tokens = tokens_info.get('prompt_tokens', 0)
        self.completion_tokens = tokens_info.get('completion_tokens', 0)
        self.token_count = self.completion_tokens
        
    def _append_message(self, sender, message, is_streaming=False):
        """添加消息到聊天窗口
        
        Args:
            sender: 消息发送者
            message: 消息内容
            is_streaming: 是否为流式响应
            
        功能:
        - 格式化消息显示样式
        - 支持Markdown渲染
        - 处理不同类型消息的显示样式
        """
        cursor = self.chat_display.textCursor()
        cursor.movePosition(cursor.End)

        if sender == self.ui_text.get('you', 'You'):
            prefix = f"\n{sender}: "
            message_html = f'<span style="color:#333333">{message}</span>'
        else:
            prefix = f"\n{self.ui_text.get('bot', 'Bot')}: "
            if is_streaming:
                message_html = f'<span style="color:#888888">{message}</span>'
            else:
                # 将Markdown转换为HTML
                try:
                    # 重置转换器状态
                    self.md.reset()
                    message_html = self.md.convert(message)
                    message_html = f'<div style="color:#2196F3">{message_html}</div>'
                except:
                    message_html = f'<span style="color:#2196F3">{message}</span>'

        cursor.insertText(prefix)
        cursor.insertHtml(message_html)
        self.chat_display.ensureCursorVisible()

    def update_bot_response(self, content):
        """更新AI响应内容
        
        Args:
            content: 新的响应内容片段
            
        功能:
        - 累积并更新AI的响应内容
        - 处理流式响应的实时显示
        - 支持Markdown实时渲染
        - 记录相关日志
        """
        # 记录AI响应，仅在第一次开始响应时记录
        current_text = self.chat_display.toPlainText()
        if self.current_response == "" and current_text.endswith(self.ui_text.get('thinking', 'Thinking...')):
            self.logger.info(self.ui_text.get('log_ai_start', 'AI starts responding...'))
            
        self.current_response += content
        cursor = self.chat_display.textCursor()
        
        if current_text.endswith(self.ui_text.get('thinking', 'Thinking...')):
            # 第一次响应，替换"思考中..."
            cursor.movePosition(cursor.End)
            for _ in range(len(self.ui_text.get('thinking', 'Thinking...'))):
                cursor.deletePreviousChar()
            # 转换并显示新的响应
            try:
                self.md.reset()
                response_html = self.md.convert(self.current_response)
                cursor.insertHtml(f'<div style="color:#2196F3">{response_html}</div>')
                self.last_cursor_position = cursor.position()
            except:
                cursor.insertHtml(f'<span style="color:#2196F3">{self.current_response}</span>')
                self.last_cursor_position = cursor.position()
        else:
            # 获取新内容的HTML
            try:
                self.md.reset()
                # 只处理新增的内容
                response_html = self.md.convert(content)
                # 将光标移动到上次的位置
                cursor.setPosition(self.last_cursor_position)
                cursor.insertHtml(f'<span style="color:#2196F3">{response_html}</span>')
                self.last_cursor_position = cursor.position()
            except:
                cursor.setPosition(self.last_cursor_position)
                cursor.insertText(content)
                self.last_cursor_position = cursor.position()
                
        self.chat_display.setTextCursor(cursor)
        self.chat_display.ensureCursorVisible()

    def finalize_response(self):
        """完成响应处理"""
        # 计算响应速度和总耗时
        elapsed_time = time.time() - self.start_time
        if elapsed_time > 0:
            tokens_per_sec = self.token_count / elapsed_time
            total_tokens = self.prompt_tokens + self.completion_tokens
            speed_text = self.ui_text.get('log_speed', 'Speed: {:.2f} Tokens/Sec | Total: {} (Input: {}, Output: {}) | Time: {:.2f}s').format(
                tokens_per_sec,
                total_tokens,
                self.prompt_tokens,
                self.completion_tokens,
                elapsed_time
            )
            self.logger.info(speed_text)
            # 更新并显示速率标签
            self.speed_label.setText(speed_text)
            self.speed_label.setVisible(True)
        
        # 记录响应完成
        self.logger.info(self.ui_text.get('log_ai_finish', 'AI response completed'))
        
        self.current_response = ""  # 清空累积的响应
        current = self.chat_display.toPlainText()
        if not current.endswith("\n"):
            self.chat_display.append("")

    def toggle_log_display(self):
        """切换日志显示状态"""
        self.log_visible = not self.log_visible
        self.log_display.setVisible(self.log_visible)

# 主程序入口
if __name__ == "__main__":
    # 设置环境变量，禁用 GLX
    os.environ["QT_QPA_PLATFORM"] = "xcb"
    os.environ["QT_XCB_GL_INTEGRATION"] = "none"
    os.environ["QT_LOGGING_RULES"] = "qt.qpa.*=false"
    
    # 必须在创建 QApplication 之前设置这些属性
    QApplication.setAttribute(Qt.AA_UseSoftwareOpenGL)
    QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps)
    QApplication.setAttribute(Qt.AA_DisableHighDpiScaling)
    
    # 创建应用程序及窗口
    app = QApplication(sys.argv)
    window = ChatWindow()
    window.show()
    sys.exit(app.exec_())